import random

from termcolor import COLORS, colored

from lagent.utils import get_logger
from .hook import Hook


class MessageLogger(Hook):
    """
    消息日志记录器，用于记录代理系统中的消息流转。
    继承自 Hook 基类，可以监听代理系统的各个阶段。
    """

    def __init__(self, name: str = 'lagent', add_file_handler: bool = False):
        """
        初始化日志记录器。

        Args:
            name: 日志器名称，默认为 'lagent'
            add_file_handler: 是否添加文件处理器，默认为 False（仅控制台输出）
        """
        # 创建日志器，配置日志格式
        self.logger = get_logger(
            name,  # 日志器名称
            'info',  # 日志级别
            '%(asctime)s %(levelname)8s %(name)8s - %(message)s',  # 日志格式
            add_file_handler=add_file_handler  # 是否添加文件处理器
        )
        # 用于存储发送者与颜色的映射关系
        self.sender2color = {}

    def before_agent(self, agent, messages, session_id):
        """
        代理处理前的钩子函数。

        Args:
            agent: 代理对象
            messages: 待处理的消息列表
            session_id: 会话ID
        """
        # 记录每条输入消息
        for message in messages:
            self._process_message(message, session_id)

    def after_agent(self, agent, message, session_id):
        """
        代理处理后的钩子函数。

        Args:
            agent: 代理对象
            message: 处理后的消息
            session_id: 会话ID
        """
        # 记录代理处理后的消息
        self._process_message(message, session_id)

    def before_action(self, executor, message, session_id):
        """
        动作执行前的钩子函数。

        Args:
            executor: 动作执行器
            message: 待执行的消息
            session_id: 会话ID
        """
        # 记录执行前的消息
        self._process_message(message, session_id)

    def after_action(self, executor, message, session_id):
        """
        动作执行后的钩子函数。

        Args:
            executor: 动作执行器
            message: 执行结果消息
            session_id: 会话ID
        """
        # 记录执行结果消息
        self._process_message(message, session_id)

    def _process_message(self, message, session_id):
        """
        处理并记录单条消息。

        Args:
            message: 要记录的消息
            session_id: 会话ID
        """
        # 获取消息发送者
        sender = message.sender
        # 为发送者分配固定的颜色（如果还未分配）
        color = self.sender2color.setdefault(sender, random.choice(list(COLORS)))
        # 使用彩色输出记录消息
        self.logger.info(colored(
            f'session id: {session_id}, message sender: {sender}\n'
            f'{message.content}', 
            color  # 使用发送者对应的颜色
        ))
